import socket
import struct
import ipaddress
import random
from subprocess import run, PIPE
from .logger import get_logger

logger = get_logger()

def parse_target_range(target):
    """该方法用于解析目标字符串，将其转换为IP地址列表 """
    try:
        if '-' in target:
            # Handle IP range (e.g., 192.168.1.1-100)
            start, end = target.rsplit('.', 1)[0], target.split('-')[1]
            base = target.rsplit('.', 1)[0]
            start_num = int(target.rsplit('.', 1)[1].split('-')[0])
            return [f"{base}.{i}" for i in range(start_num, int(end) + 1)]
        elif '/' in target:
            # Handle CIDR notation
            return [str(ip) for ip in ipaddress.ip_network(target)]
        else:
            # Single IP address
            return [target]
    except Exception as e:
        logger.error(f"解析目标范围出错: {str(e)}")
        return []

def parse_ports(ports_str, scan_type):
    """该方法用于解析端口字符串，将其转换为端口号列表"""
    try:
        ports = set()
        for part in ports_str.split(','):
            if '-' in part:
                start, end = map(int, part.split('-'))
                ports.update(range(start, end + 1))
            else:
                ports.add(int(part))
        
        # Filter UDP/TCP ports based on scan type
        if scan_type == 'udp':
            common_udp_ports = {53, 67, 68, 69, 123, 161, 162, 514, 520}
            return sorted(ports & common_udp_ports)
        return sorted(ports)
    except Exception as e:
        logger.error(f"解析端口出错: {str(e)}")
        return []

def is_host_alive(ip):
    """检查主机是否存活"""
    try:
        if run(['ping', '-n', '1', '-w', '1000', ip], stdout=PIPE, stderr=PIPE).returncode == 0:
            return True
        return False
    except Exception:
        return False

def create_syn_packet(dst_ip, dst_port):
    """创建TCP SYN包"""
    # IP Header
    ip_version = 4
    ip_ihl = 5
    ip_tos = 0
    ip_tot_len = 20 + 20  # IP header + TCP header
    ip_id = random.randint(1, 65535)
    ip_frag_off = 0
    ip_ttl = 64
    ip_proto = socket.IPPROTO_TCP
    ip_check = 0
    ip_saddr = socket.inet_aton('0.0.0.0')  # Source IP
    ip_daddr = socket.inet_aton(dst_ip)
    
    ip_ihl_version = (ip_version << 4) + ip_ihl
    ip_header = struct.pack('!BBHHHBBH4s4s',
        ip_ihl_version, ip_tos, ip_tot_len, ip_id, ip_frag_off,
        ip_ttl, ip_proto, ip_check, ip_saddr, ip_daddr)
    
    # TCP Header
    tcp_source = random.randint(1024, 65535)
    tcp_dest = dst_port
    tcp_seq = random.randint(1, 4294967295)
    tcp_ack_seq = 0
    tcp_doff = 5
    tcp_fin = 0
    tcp_syn = 1
    tcp_rst = 0
    tcp_psh = 0
    tcp_ack = 0
    tcp_urg = 0
    tcp_window = socket.htons(5840)
    tcp_check = 0
    tcp_urg_ptr = 0
    
    tcp_offset_res = (tcp_doff << 4) + 0
    tcp_flags = tcp_fin + (tcp_syn << 1) + (tcp_rst << 2) + \
                (tcp_psh << 3) + (tcp_ack << 4) + (tcp_urg << 5)
    
    tcp_header = struct.pack('!HHLLBBHHH',
        tcp_source, tcp_dest, tcp_seq, tcp_ack_seq,
        tcp_offset_res, tcp_flags, tcp_window, tcp_check, tcp_urg_ptr)
    
    return ip_header + tcp_header

def send_udp_probe(target, port, timeout):
    """发送UDP探测包并判断端口是否开放"""
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(timeout)
    
    try:
        sock.sendto(b'', (target, port))
        sock.recvfrom(1024)
        return True
    except socket.error:
        return False
    finally:
        sock.close() 